import os
import librosa
import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
from sklearn.model_selection import train_test_split
import torchaudio
from . import batch_size, test_size_rate

class SpeechDataset(Dataset):
    def __init__(self, file_paths, labels, num_mfcc=13, max_len=16000):
        """
        初始化参数
        :param file_paths: 音频文件路径列表
        :param labels: 音频文件标签列表
        :param num_mfcc: 梅尔频率倒谱系数的特征数量
        :param max_len: 音频文件的采样率
        """
        self.file_paths = file_paths
        self.labels = labels
        self.num_mfcc = num_mfcc
        self.max_len = max_len

    def __len__(self):
        """
        计算并返回文件数量
        SpeechDataset.__len__() 方法是 PyTorch 中的 Dataset 类要求实现的一个方法，用于返回数据集中样本的数量。
        它是 PyTorch 数据加载器 (DataLoader) 用来确定迭代次数的关键方法。
        """
        return len(self.file_paths)

    def __getitem__(self, idx):
        file_path = self.file_paths[idx]  # 音频路径
        label = self.labels[idx]  # 音频标签

        # 加载音频文件。返回值y：音频文件的时间序列数据（一维数组，表示音频的振幅）；sr：音频文件的采样率（Hz）。
        y, sr = librosa.load(file_path, sr=self.max_len)  # sr是样本采样率

        # 裁剪或填充到 max_len
        if len(y) < self.max_len:
            # 若音频数据长度 len(y) 小于 self.max_len，用 0 在末尾填充至 self.max_len。
            y = np.pad(y, (0, max(0, self.max_len - len(y))), 'constant')
        else:
            # 若音频数据长度 len(y) 大于 self.max_len，截取前 self.max_len 的数据。
            y = y[:self.max_len]

        # 提取MFCC（梅尔频率倒谱系数）特征
        mfcc = librosa.feature.mfcc(y=y, sr=self.max_len, n_mfcc=self.num_mfcc)  # 将音频数据 y 转换为 MFCC特征
        mfcc = np.expand_dims(mfcc, axis=0)  # 增加一个维度，使其形状为 [1, num_mfcc, time_steps]，符合模型输入需求
        mfcc = torch.tensor(mfcc, dtype=torch.float32)  # 将 mfcc 转换为 PyTorch 的张量格式

        return mfcc, label


def prepare_dataloaders(data_dir, keywords, test_size=0.2):
    """准备数据集"""
    # exist_ok=True表示即使目录已存在，也不报错，直接返回
    os.makedirs(data_dir, exist_ok=True)
    # 下载数据集
    torchaudio.datasets.SPEECHCOMMANDS(root=data_dir, url='speech_commands_v0.02', download=True)
    # 重定向数据集的路径
    data_dir = os.path.join(data_dir, 'SpeechCommands/speech_commands_v0.02')

    # 制作数据集与标签
    label_map = {word: i for i, word in enumerate(keywords)}
    file_paths, labels = [], []
    for keyword in keywords:
        keyword_dir = os.path.join(data_dir, keyword)
        for file_name in os.listdir(keyword_dir):
            file_paths.append(os.path.join(keyword_dir, file_name))
            labels.append(label_map[keyword])
    train_files, val_files, train_labels, val_labels = train_test_split(file_paths, labels,
                                                                        test_size=test_size_rate, random_state=42)

    train_dataset = SpeechDataset(train_files, train_labels)
    val_dataset = SpeechDataset(val_files, val_labels)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
    return train_loader, val_loader, label_map
